# -*- coding: utf-8 -*-
"""
Created on Fri Jun 27 11:33:39 2025

@author: Administrator
"""

import arcpy
import numpy as np
import socket
import struct


def grid_client():
    output_raster_path = input("请输入输出栅格完整路径及文件名（含扩展名）: ")
    # 创建一个 TCP/IP 套接字
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # 服务器地址和端口
    server_address = ('127.0.0.1', 12340)
    # 连接到服务器
    client_socket.connect(server_address)
    print("已连接到服务器......")

    # 1-1接收数组形状，波段、行列数
    shape_bytes = client_socket.recv(12)
    band_count, rows, cols = struct.unpack('!III', shape_bytes)

    # 2-接收d_tpye2(数据类型，如：“16_BIT_SIGNED”)字符串长度和内容
    d_tpye2_len = struct.unpack('!I', client_socket.recv(4))[0]
    d_tpye2_bytes = b''
    while len(d_tpye2_bytes) < d_tpye2_len:
        d_tpye2_bytes += client_socket.recv(d_tpye2_len - len(d_tpye2_bytes))
    d_tpye2 = d_tpye2_bytes.decode('utf-8')  # 转换为字符串

    # 3-接收 noDataValue（假设为 int32，与发送端对应）
    noDataValue = struct.unpack('!i', client_socket.recv(4))[0]

    # 4-1接收数组数据长度
    length_bytes = client_socket.recv(8)
    length = struct.unpack('!Q', length_bytes)[0]
    print("开始接收数据......")

    # 4-2接收数组数据
    raster_array_bytes = b''
    while len(raster_array_bytes) < length:
        data = client_socket.recv(length - len(raster_array_bytes))
        if not data:
            break
        raster_array_bytes += data

    # 根据类型代码映射到numpy数据类型
    dtype_map = {
        '8_BIT_SIGNED': np.int8,
        '8_BIT_UNSIGNED': np.uint8,
        '16_BIT_SIGNED': np.int16,
        '16_BIT_UNSIGNED': np.uint16,
        '32_BIT_SIGNED': np.int32,
        '32_BIT_UNSIGNED': np.uint32,
        '32_BIT_FLOAT': np.float32,
        '64_BIT': np.float64
    }
    dtype = dtype_map.get(d_tpye2, np.uint8)  # 默认uint8
    print(f"接收到的数据类型: {dtype}")

    # 转换为 NumPy 数组,reshape重塑数组为三维结构：
    raster_array = np.frombuffer(raster_array_bytes, dtype=dtype).reshape((band_count, rows, cols))

    # 接收投影信息长度（关键添加）
    sr_len_data = client_socket.recv(4)
    sr_len = struct.unpack('!I', sr_len_data)[0]

    # 接收投影信息内容
    sr_data = b''
    while len(sr_data) < sr_len:
        chunk = client_socket.recv(min(4096, sr_len - len(sr_data)))
        if not chunk:
            raise RuntimeError("连接断开")
        sr_data += chunk
    sr_wkt = sr_data.decode('utf-8')  # 转换回字符串

    # 接收属性,左下xy坐标和像元大小
    attributes_bytes = client_socket.recv(24)
    x_origin, y_origin, cell_size = struct.unpack('!ddd', attributes_bytes)

    # 将 NumPy 数组转换回栅格数据
    output_raster = arcpy.NumPyArrayToRaster(raster_array.astype(np.int16),
                                             arcpy.Point(x_origin, y_origin),
                                             cell_size,value_to_nodata = noDataValue)
    # 保存输出栅格数据
    print("正在保存数据，请稍等......")
    arcpy.management.CopyRaster(output_raster, output_raster_path,
                                pixel_type=d_tpye2)
    # 设置投影
    arcpy.management.DefineProjection(output_raster_path, sr_wkt)
    print("栅格数据导出成功！")

    # 清理
    del raster_array
    del output_raster
    # 关闭连接
    client_socket.close()
